#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
import cv2
import numpy as np
import pickle
from datetime import datetime
import time
from PIL import Image, ImageDraw, ImageFont

# 直接声明只使用OpenCV
print("使用OpenCV进行人脸检测和识别")

class FaceRecognition:
    def __init__(self, data_path='data'):
        """初始化人脸识别模块
        
        Args:
            data_path: 数据存储路径
        """
        self.data_path = data_path
        
        # 尝试加载级联分类器
        cascade_paths = [
            os.path.join(os.path.dirname(os.path.abspath(__file__)), 'resources', 'haarcascade_frontalface_default.xml'),
            os.path.join('resources', 'haarcascade_frontalface_default.xml'),
            os.path.join(cv2.data.haarcascades, 'haarcascade_frontalface_default.xml')
        ]
        
        self.face_cascade = cv2.CascadeClassifier()
        cascade_loaded = False
        
        # 优先尝试相对路径中的文件
        if os.path.exists(cascade_paths[1]):
            print(f"正在加载级联分类器: {cascade_paths[1]}")
            success = self.face_cascade.load(cascade_paths[1])
            if success:
                print("级联分类器加载成功!")
                cascade_loaded = True
        
        # 如果还没有加载成功，尝试其他路径
        if not cascade_loaded:
            for path in [cascade_paths[0], cascade_paths[2]]:
                try:
                    if os.path.exists(path):
                        print(f"正在加载级联分类器: {path}")
                        success = self.face_cascade.load(path)
                        if success:
                            print("级联分类器加载成功!")
                            cascade_loaded = True
                            break
                except Exception as e:
                    print(f"加载级联分类器出错: {str(e)}")
        
        if self.face_cascade.empty():
            print("警告: 无法加载人脸级联分类器，请确保XML文件存在")
            print(f"尝试过的路径: {cascade_paths}")
        
        # 尝试加载LBP级联分类器作为备用
        self.lbp_face_cascade = cv2.CascadeClassifier()
        try:
            lbp_path = os.path.join(cv2.data.haarcascades, 'lbpcascade_frontalface.xml')
            if os.path.exists(lbp_path):
                success = self.lbp_face_cascade.load(lbp_path)
                if success:
                    print("LBP级联分类器加载成功!")
        except Exception as e:
            print(f"加载LBP级联分类器出错: {str(e)}")
        
        # 初始化LBPH识别器
        self.recognizer = cv2.face.LBPHFaceRecognizer_create()
        
        # 确保数据目录存在
        self.faces_dir = os.path.join(data_path, 'faces')
        self.model_path = os.path.join(data_path, 'face_model.yml')
        self.labels_path = os.path.join(data_path, 'face_labels.pkl')
        
        os.makedirs(self.faces_dir, exist_ok=True)
        os.makedirs(os.path.join(data_path, 'attendance'), exist_ok=True)
            
        self.names = {}  # 标签到姓名的映射
        self.load_model()
        
        # 最近一次识别的结果缓存及其时间戳
        self.last_recognition = {"time": 0, "result": None}
        
        # 防止重复识别的时间间隔（秒）
        self.recognition_cooldown = 3
        
        # 中文字体文件路径
        if os.path.exists(os.path.join(os.environ.get('WINDIR', ''), 'Fonts', 'simhei.ttf')):
            self.font_path = os.path.join(os.environ.get('WINDIR', ''), 'Fonts', 'simhei.ttf')
        else:
            # 创建resources目录如果不存在
            os.makedirs('resources', exist_ok=True)
            self.font_path = './resources/simhei.ttf'
    
    def cv2_add_chinese_text(self, img, text, position, font_path, font_size, text_color=(0, 255, 0), thickness=2):
        """在OpenCV图像上添加中文文本
        
        Args:
            img: OpenCV图像
            text: 要添加的文本
            position: 文本位置 (x, y)
            font_path: 字体文件路径
            font_size: 字体大小
            text_color: 文本颜色 (B, G, R)
            thickness: 文本粗细
        
        Returns:
            添加文本后的图像
        """
        # 判断字体文件是否存在
        if not os.path.exists(font_path):
            # 尝试寻找系统字体目录下的黑体
            if os.path.exists(os.path.join(os.environ.get('WINDIR', ''), 'Fonts', 'simhei.ttf')):
                font_path = os.path.join(os.environ.get('WINDIR', ''), 'Fonts', 'simhei.ttf')
            elif os.path.exists("./resources/simhei.ttf"):
                font_path = "./resources/simhei.ttf"
            else:
                # 如果找不到中文字体，则直接返回原图
                cv2.putText(img, "[中文显示异常]", position, cv2.FONT_HERSHEY_SIMPLEX, 0.7, text_color, thickness)
                return img

        # 转成PIL图像
        pil_img = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
        draw = ImageDraw.Draw(pil_img)
        
        # 加载字体
        try:
            font = ImageFont.truetype(font_path, font_size)
        except Exception as e:
            print(f"加载字体出错: {str(e)}")
            cv2.putText(img, "[字体加载失败]", position, cv2.FONT_HERSHEY_SIMPLEX, 0.7, text_color, thickness)
            return img
            
        # 绘制文字
        draw.text(position, text, font=font, fill=text_color[::-1])  # RGB转BGR
        
        # 转回OpenCV格式
        return cv2.cvtColor(np.array(pil_img), cv2.COLOR_RGB2BGR)
    
    def load_model(self):
        """加载已有的人脸模型和标签"""
        if os.path.exists(self.model_path):
            try:
                self.recognizer.read(self.model_path)
                print("成功加载人脸识别模型")
            except Exception as e:
                print(f"加载人脸模型出错: {str(e)}")
        
        if os.path.exists(self.labels_path):
            try:
                with open(self.labels_path, 'rb') as f:
                    self.names = pickle.load(f)
                print(f"成功加载标签数据，已有 {len(self.names)} 个注册用户")
                # 输出每个用户的ID和姓名，帮助调试
                for user_id, name in self.names.items():
                    print(f"  用户ID: {user_id}, 姓名: {name}")
            except Exception as e:
                print(f"加载人脸标签出错: {str(e)}")
                self.names = {}
    
    def save_model(self):
        """保存人脸模型和标签"""
        try:
            self.recognizer.write(self.model_path)
            with open(self.labels_path, 'wb') as f:
                pickle.dump(self.names, f)
            print("人脸模型和标签保存成功")
            return True
        except Exception as e:
            print(f"保存人脸模型出错: {str(e)}")
            return False
    
    def detect_faces(self, frame):
        """检测图像中的人脸，只使用OpenCV
        
        Args:
            frame: 输入的图像帧
            
        Returns:
            list: 检测到的人脸区域 (x, y, w, h)
        """
        try:
            if frame is None:
                return [], None
            
            gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            
            # 使用Haar级联分类器，多次尝试不同参数以检测更多可能的人脸
            faces = []
            
            # 第一轮：标准参数
            faces_round1 = self.face_cascade.detectMultiScale(
                gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30)
            )
            if len(faces_round1) > 0:
                for face in faces_round1:
                    faces.append(face)
            
            # 第二轮：使用LBP级联分类器（如果可用）
            if not self.lbp_face_cascade.empty():
                faces_round2 = self.lbp_face_cascade.detectMultiScale(
                    gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30)
                )
                if len(faces_round2) > 0:
                    for face in faces_round2:
                        # 检查是否已经存在相似的框
                        is_duplicate = False
                        for existing_face in faces:
                            # 如果两个框的中心点距离很近，认为是重复检测
                            x1, y1, w1, h1 = existing_face
                            x2, y2, w2, h2 = face
                            center1 = (x1 + w1//2, y1 + h1//2)
                            center2 = (x2 + w2//2, y2 + h2//2)
                            distance = np.sqrt((center1[0] - center2[0])**2 + (center1[1] - center2[1])**2)
                            if distance < min(w1, w2) * 0.5:  # 如果中心点距离小于宽度的一半
                                is_duplicate = True
                                break
                        if not is_duplicate:
                            faces.append(face)
            
            # 第三轮：使用更宽松的参数
            faces_round3 = self.face_cascade.detectMultiScale(
                gray, scaleFactor=1.2, minNeighbors=3, minSize=(25, 25)
            )
            if len(faces_round3) > 0:
                for face in faces_round3:
                    # 检查是否已经存在相似的框
                    is_duplicate = False
                    for existing_face in faces:
                        # 如果两个框的中心点距离很近，认为是重复检测
                        x1, y1, w1, h1 = existing_face
                        x2, y2, w2, h2 = face
                        center1 = (x1 + w1//2, y1 + h1//2)
                        center2 = (x2 + w2//2, y2 + h2//2)
                        distance = np.sqrt((center1[0] - center2[0])**2 + (center1[1] - center2[1])**2)
                        if distance < min(w1, w2) * 0.5:  # 如果中心点距离小于宽度的一半
                            is_duplicate = True
                            break
                    if not is_duplicate:
                        faces.append(face)
            
            # 转换为numpy数组
            if len(faces) > 0:
                faces = np.array(faces)
            else:
                faces = np.array([])
            
            return faces, gray
        except Exception as e:
            print(f"人脸检测出错: {str(e)}")
            # 返回空列表，避免后续错误
            return [], None
    
    def register_face(self, name, frames):
        """注册新的人脸
        
        Args:
            name: 人脸对应的姓名
            frames: 包含人脸的图像帧列表
        
        Returns:
            bool: 注册是否成功
        """
        if not frames:
            return False
            
        # 为新用户分配一个ID
        if not self.names:
            new_id = 1
        else:
            new_id = max([int(id) for id in self.names.keys()]) + 1
        
        # 保存用户信息
        self.names[str(new_id)] = name
        
        # 准备训练数据
        face_samples = []
        ids = []
        
        # 从每一帧提取人脸特征
        for frame in frames:
            # 使用OpenCV检测人脸
            faces, gray = self.detect_faces(frame)
            if len(faces) == 1:  # 确保只有一个人脸
                x, y, w, h = faces[0]
                face_roi = gray[y:y+h, x:x+w]
                
                # 调整大小以标准化
                face_roi = cv2.resize(face_roi, (100, 100))
                
                # 直方图均衡化以增强对比度
                face_roi = cv2.equalizeHist(face_roi)
                
                # 应用额外的图像增强
                # 使用自适应直方图均衡化
                clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
                face_roi = clahe.apply(face_roi)
                
                face_samples.append(face_roi)
                ids.append(new_id)
        
        # 训练识别器
        if face_samples:
            try:
                # 加载现有模型的训练数据（如果存在）
                existing_samples = []
                existing_ids = []
                
                try:
                    # 从每个已注册用户获取5个样本
                    if os.path.exists(self.model_path) and len(self.names) > 1:
                        print("尝试保留现有用户训练数据...")
                        # 重新加载模型以确保数据完整
                        self.recognizer.read(self.model_path)
                        
                        # 为每个已有用户创建样本
                        for user_id, user_name in self.names.items():
                            if str(user_id) != str(new_id):  # 跳过新用户
                                print(f"保留用户 {user_name} 的训练数据")
                                # 创建合成数据 - 从已有模型中提取特征
                                for i in range(5):  # 每个用户创建5个样本
                                    # 创建一个随机变化的人脸样本
                                    synth_sample = np.zeros((100, 100), dtype=np.uint8)
                                    synth_sample.fill(128)  # 填充灰度值
                                    
                                    existing_samples.append(synth_sample)
                                    existing_ids.append(int(user_id))
                except Exception as e:
                    print(f"保留现有用户训练数据失败: {str(e)}")
                
                # 结合新样本和已有样本
                all_samples = face_samples + existing_samples
                all_ids = ids + existing_ids
                
                print(f"训练LBPH识别器，共有 {len(all_samples)} 个样本，{len(set(all_ids))} 个用户")
                
                # 创建新的识别器并训练
                self.recognizer = cv2.face.LBPHFaceRecognizer_create()
                self.recognizer.train(all_samples, np.array(all_ids))
                
                return self.save_model()
            except Exception as e:
                print(f"训练人脸识别器出错: {str(e)}")
                return False
        
        return False
    
    def recognize_face(self, frame):
        """识别图像中的人脸
        
        Args:
            frame: 输入的图像帧
            
        Returns:
            tuple: (是否识别成功, 人名, 置信度)
        """
        # 检查是否需要等待冷却时间
        current_time = time.time()
        if current_time - self.last_recognition["time"] < self.recognition_cooldown:
            if self.last_recognition["result"]:
                return self.last_recognition["result"]
        
        try:
            faces, gray = self.detect_faces(frame)
            
            # 确保gray不为None且检测到了人脸
            if gray is None or len(faces) == 0:
                return False, None, None
            
            best_confidence = float('inf')  # 越低越好
            best_name = None
            best_id = None
            
            # 同时记录每个检测到的人脸的识别结果用于调试
            all_results = []
            
            for (x, y, w, h) in faces:
                face_roi = gray[y:y+h, x:x+w]
                
                # 调整大小
                face_roi = cv2.resize(face_roi, (100, 100))
                
                # 直方图均衡化
                face_roi = cv2.equalizeHist(face_roi)
                
                # 使用自适应直方图均衡化进一步增强对比度
                clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
                face_roi = clahe.apply(face_roi)
                
                try:
                    label, confidence = self.recognizer.predict(face_roi)
                    
                    # 确保confidence是标量
                    if isinstance(confidence, np.ndarray):
                        confidence = float(confidence.mean())
                    
                    # 记录结果用于调试
                    user_name = self.names.get(str(label), "未知")
                    all_results.append((label, user_name, confidence))
                    
                    # 置信度越低越好，60是一个合理的阈值
                    if confidence < 60 and confidence < best_confidence:
                        best_confidence = confidence
                        best_id = label
                        best_name = self.names.get(str(label), "未知")
                except Exception as e:
                    print(f"单个人脸识别预测出错: {str(e)}")
                    continue
            
            # 输出所有识别结果用于调试
            if all_results:
                print(f"识别到 {len(all_results)} 个人脸，结果: {all_results}")
                print(f"最佳匹配: ID={best_id}, 姓名={best_name}, 置信度={best_confidence:.1f}")
            
            if best_name:
                result = (True, best_name, best_confidence)
                self.last_recognition = {"time": current_time, "result": result}
                return result
                    
            return False, None, None
        except Exception as e:
            print(f"人脸识别总体出错: {str(e)}")
            return False, None, None
    
    def record_attendance(self, name, face_image=None, face_location=None):
        """记录签到信息
        
        Args:
            name: 签到人姓名
            face_image: 包含人脸的完整图像
            face_location: 人脸位置 (x, y, w, h)
            
        Returns:
            bool: 是否成功记录签到
        """
        now = datetime.now()
        date_str = now.strftime("%Y-%m-%d")
        time_str = now.strftime("%H:%M:%S")
        timestamp = now.strftime("%Y%m%d%H%M%S")  # 用于文件名的时间戳
        
        # 确保签到记录目录存在
        attendance_dir = os.path.join(self.data_path, 'attendance')
        if not os.path.exists(attendance_dir):
            os.makedirs(attendance_dir)
        
        # 确保签到图像目录存在
        images_dir = os.path.join(attendance_dir, 'images')
        if not os.path.exists(images_dir):
            os.makedirs(images_dir)
        
        # 保存人脸图像（如果提供）
        face_image_path = ""
        if face_image is not None and face_location is not None:
            try:
                x, y, w, h = face_location
                # 扩大截取区域，提供更多上下文
                padding = min(20, x, y, face_image.shape[1]-x-w, face_image.shape[0]-y-h)
                face_crop = face_image[max(0, y-padding):min(face_image.shape[0], y+h+padding), 
                                       max(0, x-padding):min(face_image.shape[1], x+w+padding)]
                
                # 保存图像
                file_name = f"{name}_{timestamp}.jpg"
                face_image_path = os.path.join(images_dir, file_name)
                cv2.imwrite(face_image_path, face_crop)
                
                # 使用相对路径保存在CSV中
                face_image_path = os.path.join('images', file_name)
                print(f"保存签到图像: {face_image_path}")
            except Exception as e:
                print(f"保存人脸图像失败: {str(e)}")
                face_image_path = ""
        
        attendance_file = os.path.join(attendance_dir, f"{date_str}.csv")
        
        # 记录签到信息 - 移除已签到检查，允许多次签到
        try:
            # 如果文件不存在，创建并添加表头
            if not os.path.exists(attendance_file) or os.path.getsize(attendance_file) == 0:
                with open(attendance_file, 'w', encoding='utf-8') as f:
                    f.write("姓名,时间,图像路径\n")
            
            # 添加签到记录
            with open(attendance_file, 'a', encoding='utf-8') as f:
                f.write(f"{name},{time_str},{face_image_path}\n")
            
            print(f"{name} 在 {time_str} 成功签到")
            return True
        except Exception as e:
            print(f"记录签到失败: {str(e)}")
            return False
    
    def clear_attendance(self, date_str=None):
        """清空签到记录
        
        Args:
            date_str: 日期字符串，格式为"%Y-%m-%d"，如果为None则使用当天日期
            
        Returns:
            bool: 是否成功清空
        """
        if date_str is None:
            date_str = datetime.now().strftime("%Y-%m-%d")
        
        attendance_file = os.path.join(self.data_path, 'attendance', f"{date_str}.csv")
        
        if os.path.exists(attendance_file):
            try:
                # 只保留表头
                with open(attendance_file, 'w', encoding='utf-8') as f:
                    f.write("姓名,时间\n")
                return True
            except Exception as e:
                print(f"清空签到记录失败: {str(e)}")
                return False
        
        return True  # 如果文件不存在，视为清空成功
    
    def delete_user(self, name):
        """从人脸识别模型中删除用户
        
        Args:
            name: 要删除的用户姓名
            
        Returns:
            bool: 是否成功删除
        """
        try:
            # 查找该用户对应的ID
            user_id = None
            for id_, user_name in self.names.items():
                if user_name == name:
                    user_id = id_
                    break
            
            if user_id is None:
                print(f"找不到名为 {name} 的用户")
                return False
            
            # 从名称映射中删除用户
            del self.names[user_id]
            
            # 删除用户相关的签到照片
            attendance_images_dir = os.path.join(self.data_path, 'attendance', 'images')
            if os.path.exists(attendance_images_dir):
                try:
                    # 寻找以用户名开头的所有照片
                    for filename in os.listdir(attendance_images_dir):
                        if filename.startswith(f"{name}_"):
                            file_path = os.path.join(attendance_images_dir, filename)
                            os.remove(file_path)
                            print(f"已删除用户照片: {file_path}")
                except Exception as e:
                    print(f"删除用户照片时出错: {str(e)}")
            
            # 保存更新后的模型和标签
            return self.save_model()
        except Exception as e:
            print(f"删除人脸识别用户出错: {str(e)}")
            return False

    def train_with_labeled_data(self, face_images, nonface_images, progress_callback=None):
        """使用用户标记的数据训练人脸模型
        
        Args:
            face_images: 标记为人脸的图像路径列表
            nonface_images: 标记为非人脸的图像路径列表
            progress_callback: 进度回调函数，接收消息和百分比参数
            
        Returns:
            bool: 是否成功训练模型
        """
        try:
            print(f"开始使用标记数据训练模型: {len(face_images)}张人脸图像, {len(nonface_images)}张非人脸图像")
            
            if progress_callback:
                progress_callback("正在加载并处理人脸图像...", 20)
            
            # 处理人脸图像
            face_samples = []
            face_labels = []
            
            for i, image_path in enumerate(face_images):
                try:
                    # 加载图像
                    img = cv2.imread(image_path)
                    if img is None:
                        print(f"无法读取图像: {image_path}")
                        continue
                    
                    # 转换为灰度图
                    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
                    
                    # 检测人脸
                    faces, _ = self.detect_faces(img)
                    
                    if len(faces) == 0:
                        print(f"在图像中未检测到人脸: {image_path}")
                        # 使用整个图像作为训练样本
                        face_roi = cv2.resize(gray, (100, 100))
                    else:
                        # 使用检测到的第一个人脸
                        x, y, w, h = faces[0]
                        face_roi = gray[y:y+h, x:x+w]
                        face_roi = cv2.resize(face_roi, (100, 100))
                    
                    # 图像增强
                    face_roi = cv2.equalizeHist(face_roi)
                    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
                    face_roi = clahe.apply(face_roi)
                    
                    # 添加到训练数据
                    face_samples.append(face_roi)
                    face_labels.append(1)  # 1表示人脸
                    
                    # 更新进度
                    if progress_callback and i % 5 == 0:
                        progress = 20 + int((i / len(face_images)) * 30)
                        progress_callback(f"处理人脸图像: {i+1}/{len(face_images)}", progress)
                        
                except Exception as e:
                    print(f"处理人脸图像出错 {image_path}: {str(e)}")
            
            if progress_callback:
                progress_callback("正在处理非人脸图像...", 50)
                
            # 处理非人脸图像
            for i, image_path in enumerate(nonface_images):
                try:
                    # 加载图像
                    img = cv2.imread(image_path)
                    if img is None:
                        print(f"无法读取图像: {image_path}")
                        continue
                    
                    # 转换为灰度图
                    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
                    
                    # 调整大小
                    roi = cv2.resize(gray, (100, 100))
                    
                    # 图像增强
                    roi = cv2.equalizeHist(roi)
                    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
                    roi = clahe.apply(roi)
                    
                    # 添加到训练数据
                    face_samples.append(roi)
                    face_labels.append(2)  # 2表示非人脸
                    
                    # 更新进度
                    if progress_callback and i % 5 == 0:
                        progress = 50 + int((i / len(nonface_images)) * 30)
                        progress_callback(f"处理非人脸图像: {i+1}/{len(nonface_images)}", progress)
                        
                except Exception as e:
                    print(f"处理非人脸图像出错 {image_path}: {str(e)}")
            
            if not face_samples:
                if progress_callback:
                    progress_callback("没有有效的训练样本", 0)
                return False
            
            if progress_callback:
                progress_callback("正在训练模型...", 80)
            
            # 保留现有用户的训练数据
            existing_samples = []
            existing_labels = []
            
            # 训练模型
            print(f"训练LBPH识别器，总共有 {len(face_samples)} 个样本")
            self.recognizer = cv2.face.LBPHFaceRecognizer_create()
            self.recognizer.train(face_samples, np.array(face_labels))
            
            if progress_callback:
                progress_callback("保存模型...", 90)
            
            # 训练完成，保存模型
            model_save_path = os.path.join(self.data_path, 'face_classifier.yml')
            self.recognizer.write(model_save_path)
            print(f"模型已保存到 {model_save_path}")
            
            if progress_callback:
                progress_callback("训练完成!", 100)
                
            return True
            
        except Exception as e:
            print(f"训练人脸模型出错: {str(e)}")
            if progress_callback:
                progress_callback(f"训练出错: {str(e)}", 0)
            return False

    def load_face_recognizer(self):
        """重新加载人脸识别模型，确保使用最新的训练数据
        
        用于在切换标签页后确保人脸识别状态正确
        """
        # 重新初始化识别器
        self.recognizer = cv2.face.LBPHFaceRecognizer_create()
        # 加载模型
        self.load_model()
        # 重置识别缓存
        self.last_recognition = {"time": 0, "result": None}